import os
import argparse
from tqdm import tqdm
import numpy as np
from pprint import pprint
import matplotlib.pyplot as plt

from data.dataset import build_dataset

def parse_args():
    parser = argparse.ArgumentParser(description='analysis boxes')
    parser.add_argument("--test-datasets", type=str, help='datasets name', required=True)
    parser.add_argument("--area-thresholds", type=str, help='thresholds for area', required=True)
    args = parser.parse_args()
    return args

def analysis_boxes_area_distribution(dataset_dicts, thresholds):
    def box_generator(dataset_dicts):
        for dataset_dict in dataset_dicts:
            for annotation in dataset_dict["annotations"]:
                yield annotation["bbox"]

    def fall_into_one_scope(value, thresholds):
        for i, threshold in enumerate(thresholds):
            if value < threshold:
                return i
        return len(thresholds)

    assert len(thresholds) > 0

    thresholds.sort()

    count_list = [0 for i in range(len(thresholds) + 1)]

    for box in box_generator(dataset_dicts):
        area = (box[3] - box[1]) * (box[2] - box[0])
        count_list[fall_into_one_scope(area, thresholds)] += 1

    analysis_dict = {
        "num_boxes": sum(count_list),
    }
    for lower, upper, value in zip([None] + thresholds, thresholds + [None], count_list):
        if lower is None:
            label = "0 <= area < {}".format(upper)
        elif upper is None:
            label = "{} <= area < inf".format(lower)
        else:
            label = "{} <= area < {}".format(lower, upper)
        analysis_dict[label] = value

    return analysis_dict

def analysis_boxes_count_distribution(dataset_dicts, output_dir):
    all_boxes_count_list = [len(dataset_dict['annotations']) for dataset_dict in dataset_dicts]
    class_boxes_count = {}
    for dataset_dict in dataset_dicts:
        count_dict = {}
        for obj in dataset_dict['annotations']:
            if obj['category_id'] not in count_dict:
                count_dict[obj['category_id']] = 0
            count_dict[obj['category_id']] += 1
        for class_id, count in count_dict.items():
            if class_id not in class_boxes_count:
                class_boxes_count[class_id] = []
            class_boxes_count[class_id].append(count)

    for class_id, count_list in class_boxes_count.items():
        plt.figure()
        plt.hist(count_list, bins=8, color='b', label='class {}'.format(class_id))
        plt.xlabel('boxes count')
        plt.ylabel('freq')
        plt.legend()
        output_path = os.path.join(output_dir, 'class-{}-boxes-count.png'.format(class_id))
        plt.savefig(output_path, dpi=240)

    plt.figure()
    plt.hist(all_boxes_count_list, bins=8, color='b', label='all class')
    plt.xlabel('boxes count')
    plt.ylabel('freq')
    plt.legend()
    output_path = os.path.join(output_dir, 'all-class-boxes-count.png')
    plt.savefig(output_path, dpi=240)

    all_boxes_count_mean = sum(all_boxes_count_list) / len(all_boxes_count_list)
    class_boxes_count_mean = {
        class_id : sum(count_list) / len(count_list)
        for class_id, count_list in class_boxes_count.items()
    }
    class_boxes_count_mean['all'] = all_boxes_count_mean
    
    return all_boxes_count_list, class_boxes_count, class_boxes_count_mean

if __name__ == '__main__':
    args = parse_args()

    test_datasets = ["{}-test".format(dataset) for dataset in args.test_datasets.split(',')]
    dataset_dicts = build_dataset(*test_datasets)

    area_thresholds = [float(i) for i in args.area_thresholds.split(',')]
    analysis_dict = analysis_boxes_area_distribution(dataset_dicts, area_thresholds)

    output_dir = os.path.join('output', 'analysis_dataset_boxes', '&'.join(test_datasets))
    os.makedirs(output_dir, exist_ok=True)

    _, _, class_boxes_count_mean = analysis_boxes_count_distribution(dataset_dicts, output_dir)

    output_path = os.path.join(output_dir, 'area_threshold.txt')
    with open(output_path, 'w', encoding='utf-8') as f:
        pprint(analysis_dict, stream=f)
        pprint(class_boxes_count_mean, stream=f)

    print("done")
        
    

